#include "YOLODetector.h"
#include <iostream>

/**
 * 构造函数 - 初始化基本参数
 */
YOLODetector::YOLODetector() {
    // 设置模型输入尺寸为640x640
    inputSize = cv::Size(640, 640);
    
    // 初始化分类类别名称
    classNames = {
        "Default", "Class_1", "Class_2", "Class_3", "Class_4", "Class_5", 
        "Class_6", "Class_7", "Class_8", "Class_9", "Class_10", 
        "Class_11", "Class_12", "Class_13"
    };
}

/**
 * 初始化模型 - 加载ONNX模型并配置推理参数
 * 
 * @param modelPath ONNX模型文件路径
 * @return 是否成功初始化模型
 */
bool YOLODetector::initialize(const std::string& modelPath) {
    try {
        // 加载ONNX模型
        net = cv::dnn::readNetFromONNX(modelPath);
        
        // 设置推理后端和目标设备
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
        
        std::cout << "模型加载成功: " << modelPath << std::endl;
        return true;
    } catch (const cv::Exception& e) {
        std::cerr << "加载模型失败: " << e.what() << std::endl;
        return false;
    }
}

/**
 * 图像预处理 - 调整大小并进行归一化
 * 
 * @param image 输入原始图像
 * @return 处理后的Blob数据，适合神经网络输入
 */
cv::Mat YOLODetector::preprocess(const cv::Mat& image) {
    // 创建神经网络输入Blob
    cv::Mat blob;
    
    // 图像预处理：调整大小、归一化和通道转换
    cv::dnn::blobFromImage(
        image,              // 输入图像
        blob,               // 输出Blob
        1.0/255.0,          // 归一化系数
        inputSize,          // 调整后的尺寸
        cv::Scalar(0,0,0),  // 均值
        true,               // 是否交换RGB通道
        false,              // 是否为单通道图像
        CV_32F              // 输出数据类型
    );
    
    return blob;
}

/**
 * 目标检测主函数 - 执行推理并获取检测结果
 * 
 * @param image 输入图像
 * @param confThreshold 置信度阈值
 * @param nmsThreshold NMS阈值
 * @return 检测结果列表
 */
std::vector<Detection> YOLODetector::detect(const cv::Mat& image, float confThreshold, float nmsThreshold) {
    // 验证输入图像有效性
    if (image.empty()) {
        std::cerr << "输入图像为空!" << std::endl;
        return {};
    }
    
    // 1. 图像预处理
    cv::Mat blob = preprocess(image);
    
    // 2. 设置网络输入
    net.setInput(blob);
    
    // 3. 获取输出层名称
    std::vector<std::string> outLayerNames = net.getUnconnectedOutLayersNames();
    
    // 4. 执行前向推理
    std::vector<cv::Mat> outputs;
    net.forward(outputs, outLayerNames);
    
    // 5. 输出基本信息用于调试
    std::cout << "模型输出数量: " << outputs.size() << std::endl;
    if (!outputs.empty()) {
        std::cout << "输出维度: ";
        for (int d = 0; d < outputs[0].dims; d++) {
            std::cout << outputs[0].size[d] << " ";
        }
        std::cout << std::endl;
    }
    
    // 6. 后处理并返回检测结果
    return postprocess(outputs[0], image.size(), confThreshold, nmsThreshold);
}

/**
 * 后处理函数 - 处理网络输出并转换为检测结果
 * 
 * @param output 网络输出数据
 * @param originalImageSize 原始图像尺寸
 * @param confThreshold 置信度阈值
 * @param nmsThreshold NMS阈值
 * @return 检测结果列表
 */
std::vector<Detection> YOLODetector::postprocess(const cv::Mat& output, const cv::Size& originalImageSize, 
                                               float confThreshold, float nmsThreshold) {
    std::vector<Detection> detections;
    
    // 输出原始模型输出形状信息
    std::cout << "输入形状: [";
    for (int i = 0; i < output.dims; ++i) {
        std::cout << output.size[i] << " ";
    }
    std::cout << "]" << std::endl;
    
    // 验证输入数据维度
    if (output.dims < 3) {
        std::cerr << "输入维度不足，期望3维数据" << std::endl;
        return detections;
    }
    
    // 处理YOLO输出格式
    cv::Mat output_reshaped;
    
    try {
        // 处理YOLOv8标准输出格式 [1, 18, 8400]
        if (output.dims == 3) {
            std::cout << "检测到YOLO格式输出 [1, " << output.size[1] << ", " << output.size[2] << "]" << std::endl;
            
            // 直接使用原始输出格式
            output_reshaped = output;
        } else {
            std::cerr << "不支持的输出格式" << std::endl;
            return detections;
        }
    } catch (const cv::Exception& e) {
        std::cerr << "矩阵操作错误: " << e.what() << std::endl;
        return detections;
    }
    
    // 计算图像缩放比例，用于将检测框映射回原图
    float xScale = static_cast<float>(originalImageSize.width) / inputSize.width;
    float yScale = static_cast<float>(originalImageSize.height) / inputSize.height;
    
    // 存储检测结果的容器
    std::vector<int> classIds;
    std::vector<float> confidences;
    std::vector<cv::Rect> boxes;
    
    // 验证输出维度
    if (output_reshaped.dims != 3) {
        std::cerr << "输出维度不正确，期望3维" << std::endl;
        return detections;
    }
    
    // 解析模型输出
    int num_classes = output_reshaped.size[1] - 4;  // 类别数量 = 总行数 - 4个坐标值
    int num_anchors = output_reshaped.size[2];      // 锚点数量（通常是8400）
    
    std::cout << "检测到 " << num_anchors << " 个锚点，" << num_classes << " 个类别" << std::endl;
    
    // 遍历所有锚点（检测结果），YOLOv8输出格式为[1, classes+4, anchors]
    for (int col = 0; col < num_anchors; col++) {
        // 提取当前锚点的所有类别概率
        std::vector<float> class_scores(num_classes);
        for (int cls = 0; cls < num_classes; cls++) {
            // 索引计算：batch=0, 行=类别索引+4(偏移), 列=当前锚点
            class_scores[cls] = output_reshaped.ptr<float>(0)[output_reshaped.size[2] * (cls + 4) + col];
        }
        
        // 找出最高置信度及其对应的类别
        int classIdPoint = 0;
        float maxScore = 0;
        for (int cls = 0; cls < num_classes; cls++) {
            if (class_scores[cls] > maxScore) {
                maxScore = class_scores[cls];
                classIdPoint = cls;
            }
        }
        
        // 筛选高于阈值的检测结果
        if (maxScore > confThreshold) {
            try {
                // 获取边界框坐标 (XYWH格式)
                float x = output_reshaped.ptr<float>(0)[output_reshaped.size[2] * 0 + col];  // 中心点x坐标
                float y = output_reshaped.ptr<float>(0)[output_reshaped.size[2] * 1 + col];  // 中心点y坐标
                float w = output_reshaped.ptr<float>(0)[output_reshaped.size[2] * 2 + col];  // 宽度
                float h = output_reshaped.ptr<float>(0)[output_reshaped.size[2] * 3 + col];  // 高度
                
                // 转换为左上角坐标并缩放到原图大小
                int left = static_cast<int>((x - w/2) * xScale);
                int top = static_cast<int>((y - h/2) * yScale);
                int width = static_cast<int>(w * xScale);
                int height = static_cast<int>(h * yScale);
                
                // 保存边界框
                boxes.push_back(cv::Rect(left, top, width, height));
                
                // 验证并保存类别ID
                int classId = classIdPoint;
                if (classId >= 0 && classId < static_cast<int>(classNames.size())) {
                    classIds.push_back(classId);
                } else {
                    classIds.push_back(0);  // 使用默认类别
                }
                
                // 保存置信度
                confidences.push_back(maxScore);
            } catch (const cv::Exception& e) {
                std::cerr << "处理第 " << col << " 列时发生错误: " << e.what() << std::endl;
            }
        }
    }
    
    // 应用非极大值抑制(NMS)，消除重叠框
    std::vector<int> indices;
    if (!boxes.empty()) {
        cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
    }
    
    // 整合最终检测结果
    for (size_t i = 0; i < indices.size(); i++) {
        int idx = indices[i];
        Detection detection;
        detection.class_id = classIds[idx];
        detection.confidence = confidences[idx];
        detection.box = boxes[idx];
        detection.className = classNames[detection.class_id];
        
        detections.push_back(detection);
    }
    
    return detections;
}

/**
 * 可视化函数 - 在图像上绘制检测结果
 * 
 * @param image 输入/输出图像
 * @param detections 检测结果列表
 */
void YOLODetector::drawDetections(cv::Mat& image, const std::vector<Detection>& detections) {
    for (const auto& detection : detections) {
        // 根据类别ID生成一致的颜色
        cv::Scalar color(0, 255 * (detection.class_id * 123 % 255) / 255.0, 255 * (detection.class_id * 231 % 255) / 255.0);
        
        // 绘制边界框
        cv::rectangle(image, detection.box, color, 2);
        
        // 准备标签文本（类别+置信度）
        std::string label = detection.className + ": " + std::to_string(static_cast<int>(detection.confidence * 100)) + "%";
        
        // 绘制标签背景
        int baseLine;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        cv::rectangle(image, 
                     cv::Point(detection.box.x, detection.box.y - labelSize.height - baseLine),
                     cv::Point(detection.box.x + labelSize.width, detection.box.y),
                     color, 
                     cv::FILLED);
                     
        // 绘制标签文本
        cv::putText(image, 
                   label, 
                   cv::Point(detection.box.x, detection.box.y - baseLine),
                   cv::FONT_HERSHEY_SIMPLEX, 
                   0.5, 
                   cv::Scalar(255, 255, 255), 
                   1);
    }
} 